import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

def compute_loss(pred,label,original_size = None,resize_size = None,distance = 1):
    # 传入pred和lable的单张图片的坐标
    ## 用来记录距离小于指定数值的正样本个数
    positive = 0
    ## 用来记录l2 loss
    l2_loss = []
    assert len(pred) == len(label),'the len of pred and label shoule be same'
    ## 输入全部的pred和label，用list中包含tuple的形式，其中像素间距为0.139mm
    for index in range(len(pred)):
        single_pred = pred[index]
        single_label = label[index]
        ## 遇到缺失的label不进行计算
        if 0 in single_label:
            continue
        ## 计算两个点之间的像素距离。注意根据原始尺寸和resize后的尺寸对坐标值进行转换。
        if original_size is not None:
            loss = (((single_pred[0]-single_label[0])/resize_size[0]*original_size[0])**2 +
                    ((single_pred[1]-single_label[1])/resize_size[1]*original_size[1])** 2)**0.5
        else:
            ## 如果使用原图上的坐标。直接进行计算即可，不需要进行转换坐标。
            loss = (((single_pred[0] - single_label[0])) ** 2 +
                    ((single_pred[1] - single_label[1])) ** 2) ** 0.5
        ## 将计算出的像素距离转换成物理距离
        distance_loss = float(loss)*0.139
        # print('distance loss',distance_loss)
        l2_loss.append(distance_loss)
        if distance_loss <= distance:
            positive +=1
    return positive,l2_loss

def category_loss(pred_csvdir, label_csvdir, distance=2.):
    ## 计算三个标记和审核者标记之间的差异l2 loss
    all_l2_loss = []
    all_positive = 0
    pred_data = pd.read_csv(pred_csvdir)
    label_data = pd.read_csv(label_csvdir)
    ##为每个节点生成一个list用来存放loss。节点名称和list为dict关系
    name_list = (label_data.columns.values.tolist())[2:]
    name_dict = {}
    for name in name_list:
        name_dict[name] = []
    for index in range(len(label_data)):
        for i in range(2, label_data.shape[1]):
            if pd.isnull(label_data.iloc[index, i]):
                continue
            pred = [[int(float(pred_data.iloc[index, i].split(',')[-3])),
                     int(float(pred_data.iloc[index, i].split(',')[-2]))]]
            label = [[int(label_data.iloc[index, i].split(',')[-3]), int(label_data.iloc[index, i].split(',')[-2])]]
            positive, l2_loss = compute_loss(pred, label, distance=distance)
            all_positive += positive
            all_l2_loss = all_l2_loss + l2_loss
            name_dict[name_list[i - 2]] += l2_loss
    
    # print('--recall', all_positive / len(all_l2_loss), 'l2 mean', np.mean(all_l2_loss), 'l2 medium',
    #       np.median(all_l2_loss))
    name_dict = {k: np.mean(v) for k, v in name_dict.items()}
    name_list_1 = set([str(i[-2:]) for i in name_dict.keys()])
    name_list_2 = set([str(i[:-2]) for i in name_dict.keys()])
    name_dict_1 = {k: [] for k in name_list_1}
    name_dict_2 = {k: [] for k in name_list_2}
    for k, v in name_dict.items():
        name_dict_1[k[-2:]].append(v)
        name_dict_2[k[:-2]].append(v)
    name_dict = sorted(name_dict.items(), key=lambda d: d[1], reverse=True)
    name_dict_1 = {k: np.mean(v) for k, v in name_dict_1.items()}
    name_dict_2 = {k: np.mean(v) for k, v in name_dict_2.items()}
    name_dict_1 = sorted(name_dict_1.items(), key=lambda d: d[1], reverse=True)
    name_dict_2 = sorted(name_dict_2.items(), key=lambda d: d[1], reverse=True)
    
    # print(name_dict_1)
    # print(name_dict_2)
    return all_positive / len(all_l2_loss)


def get_label_loss(csv_dir, distance=2.):
    ## 计算三个标记和审核者标记之间的差异l2 loss
    all_l2_loss = []
    all_positive = 0
    data = pd.read_csv(csv_dir)
    ##为每个节点生成一个list用来存放loss。节点名称和list为dict关系
    name_list = (data.columns.values.tolist())[2:]
    name_dict = {}
    for name in name_list:
        name_dict[name] = []
    ## 根据列名进行筛选，只计算部分关键点。
    # use_name = [i for i in name_list if 'FC' in i]
    # data = data.loc[:,use_name]
    for index in range(len(data)):
        for i in range(2, data.shape[1]):
            if pd.isnull(data.iloc[index, i]):
                continue
            anno = data.iloc[index, i].split(',')
            label = [[int(anno[-3]), int(anno[-2])]]
            label_1 = [[int(anno[0]), int(anno[1])]]
            label_2 = [[int(anno[3]), int(anno[4])]]
            
            positive, l2_loss = compute_loss(label_1, label, distance=distance)
            all_positive += positive
            all_l2_loss = all_l2_loss + l2_loss
            name_dict[name_list[i - 2]] += l2_loss
            
            positive, l2_loss = compute_loss(label_2, label, distance=distance)
            all_positive += positive
            all_l2_loss = all_l2_loss + l2_loss
            name_dict[name_list[i - 2]] += l2_loss
            
            if len(anno) > 9:
                label_3 = [[int(anno[6]), int(anno[7])]]
                positive, l2_loss = compute_loss(label_3, label, distance=distance)
                all_positive += positive
                all_l2_loss = all_l2_loss + l2_loss
                name_dict[name_list[i - 2]] += l2_loss
    # print('--recall', all_positive / len(all_l2_loss), 'l2 mean', np.mean(all_l2_loss), 'l2 medium',
    #       np.median(all_l2_loss))
    name_dict = {k: np.mean(v) for k, v in name_dict.items()}
    name_list_1 = set([str(i[-2:]) for i in name_dict.keys()])
    name_list_2 = set([str(i[:-2]) for i in name_dict.keys()])
    name_dict_1 = {k: [] for k in name_list_1}
    name_dict_2 = {k: [] for k in name_list_2}
    for k, v in name_dict.items():
        name_dict_1[k[-2:]].append(v)
        name_dict_2[k[:-2]].append(v)
    name_dict = sorted(name_dict.items(), key=lambda d: d[1], reverse=True)
    name_dict_1 = {k: np.mean(v) for k, v in name_dict_1.items()}
    name_dict_2 = {k: np.mean(v) for k, v in name_dict_2.items()}
    name_dict_1 = sorted(name_dict_1.items(), key=lambda d: d[1], reverse=True)
    name_dict_2 = sorted(name_dict_2.items(), key=lambda d: d[1], reverse=True)
    
    # print(name_dict_1)
    # print(name_dict_2)
    return all_positive / len(all_l2_loss)


if __name__ == '__main__':
    pred_l2_list = []
    label_l2_list = []
    for i in [0.5+0.5*i for i in range(10)]:
        
        print(i)
        l2_recall = category_loss(pred_csvdir='../data/csv_label/0522/zhengwei/pred/anno_result.csv',
                          label_csvdir='../data/csv_label/0522/zhengwei/zhengwei_transform_val.csv',
                          distance=i)
        pred_l2_list.append(l2_recall)

        l2_recall = get_label_loss(csv_dir='../data/csv_label/0522/zhengwei/zhengwei_transform.csv', distance=i)
        label_l2_list.append(l2_recall)
        
    data = np.stack((np.array(pred_l2_list),np.array(label_l2_list)),axis=1)
    print(data)
    print(data.shape)
    s = pd.DataFrame(data,columns=['Pred','GT'],index=np.array([0.5+0.5*i for i in range(10)]))
    s.plot()
    plt.show()



    ## cewei
    pred_l2_list = []
    label_l2_list = []
    for i in [0.5 + 0.5 * i for i in range(10)]:
        print(i)
        l2_recall = category_loss(pred_csvdir='../data/csv_label/0522/cewei/pred/anno_result.csv',
                                  label_csvdir='../data/csv_label/0522/cewei/cewei_transform_val.csv',
                                  distance=i)
        pred_l2_list.append(l2_recall)
    
        l2_recall = get_label_loss(csv_dir='../data/csv_label/0522/cewei/cewei_transform.csv', distance=i)
        label_l2_list.append(l2_recall)

    data = np.stack((np.array(pred_l2_list), np.array(label_l2_list)), axis=1)
    print(data)
    print(data.shape)
    s = pd.DataFrame(data, columns=['Pred', 'GT'], index=np.array([0.5 + 0.5 * i for i in range(10)]))
    s.plot()
    plt.show()

    